import datetime
import logging
import os
import threading
import time
import traceback


import requests
import uvicorn
from fastapi import FastAPI
from gradio_client import Client

from loguru import logger
import paramiko
from starlette.background import BackgroundTasks
from starlette.responses import JSONResponse

from mysqls_test.database import SessionLocal
from mysqls_test.mysqlModels import ModelConfig, CustomizeTone

# root_dir = "E:\\"
# model_path = "D:\\Retrieval-based-Voice-Conversion-WebUI-main\\assets\\weights\\"
# index_path = "D:\\Retrieval-based-Voice-Conversion-WebUI-main\\logs\\"
root_dir = "/data/job2/"
model_path = "/data/AI/Retrieval-based-Voice-Conversion-WebUI-main/assets/weights/"
index_path = "/data/AI/Retrieval-based-Voice-Conversion-WebUI-main/logs/"

target_folder = "/data/AI/RVC-FastAPI/Data/zidingyi/"
target_real_folder = "Data/zidingyi/"
model_target_file=""
index_target_file=""

app = FastAPI()
lock = threading.Lock()

# pool = redis.ConnectionPool(host='r-bp1gguodbeuezejv25pd.redis.rds.aliyuncs.com:6379',password='aly:aly123!@#')   #实现一个连接池

@app.get("/train")
async def train(
		bakground_tasks: BackgroundTasks,
		id: int,
		step:int):
	db = SessionLocal()
	customizeTone = db.query(CustomizeTone).filter(CustomizeTone.id == id).first()
	db.close()
	if customizeTone is None:
		logging.exception("数据库读取出错")

	bakground_tasks.add_task(trainTask,customizeTone,step)
	data = {'state': 0, 'msg': '推理中', 'job_id': customizeTone.id}
	return JSONResponse(content=data)

def trainTask(customizeTone,step):

	lock.acquire()
	try:
		# r = redis.Redis(connection_pool=pool)
		# redis_lock = r.get("redis_lock")
		# # 拿到锁表示有程序正在使用GPU,等待
		# while redis_lock == "close":
		# 	time.sleep(5)
		# 	print("RVC等待锁释放打开")
		# 	redis_lock = r.get("redis_lock")
		# r.set("redis_lock","close")
		logger.info("开始异步训练task:====time:{}, id:{}",datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"),customizeTone.id)
		project_name = "train_" + str(customizeTone.id) + "_" + str(1)
		data = root_dir + project_name  # 原始语音地址
		data_vocal1 = root_dir + project_name + "_vocal1"
		data_vocal2 = root_dir + project_name + "_vocal2"
		data_vocal3 = root_dir + project_name + "_vocal3"
		data_temp = root_dir + project_name + "_temp"
		data_final = root_dir + project_name + "_final"

		#####################################################
		if step <= 0:
			down_videos(customizeTone.id, customizeTone.training_source, data)  # 下载素材音频到本地
		if step <= 1:
			uvr5_vocal(customizeTone.id, "HP3_all_vocals", data, data_temp, data_vocal1)                 # 使用UVR5分离人声
		if step <= 2:
			uvr5_vocal(customizeTone.id, "HP5_only_main_vocal", data_vocal1, data_vocal2, data_temp)      # 使用UVR5去除混响
		if step <= 3:
			uvr5_vocal(customizeTone.id, "onnx_dereverb_By_FoxJoy", data_vocal2, data_temp,data_vocal3)  # 使用UVR5去除混响
		if step <= 4:
			uvr5_vocal(customizeTone.id, "VR-DeEchoDeReverb", data_vocal3, data_final, data_temp)         # 使用UVR5去除混响
		if step <= 5:
			train_preprocess(customizeTone.id, project_name,data_final)  # 处理数据
		if step <= 6:
			train_extract_f0_feature(customizeTone.id, project_name)  # 生成特征文件
		if step <= 7:
			trainModel(customizeTone.id, project_name,data_final)  # 训练RVC模型
		if step <= 8:
			list = transfer(customizeTone.id, project_name)  # 训练好的模型上传到指定服务器
			wirteDateBase(customizeTone.id, project_name,  customizeTone.common_spk, customizeTone.common_name, list)  # 写入数据库
		# r.set("redis_lock", "open")
	except Exception as e:
		traceback.print_exc()
		logger.info("异步训练task出现异常:====time:{0}, id:{1},异常原因{2}",datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"),customizeTone.id,str(e))

	finally:
		lock.release()
		logger.info("结束异步训练task:====time:{}, id:{}",datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"),customizeTone.id)




#下载素材到指定文件夹
def down_videos(id,url,data):
	print("开始下载素材：" + str(id) +":"+ url)
	try:
		file_data = download_file(url)
		path = data + os.path.sep
		if not os.path.exists(path):
			os.mkdir(path)
		input_path = path + url.split("/")[-1]
		# 保存url中音频文件到本地服务器
		save_file(file_data, input_path)
		print("下载素材完成：" + str(id))
		# wirteMgs(id, "下载素材成功==", 1)
	except Exception as e:
		traceback.print_exc()
		# wirteMgs(id,"下载素材失败=="+str(e),3)

#使用UVR5分离人声
def uvr5_vocal(id,model,data,data_vocal,data_temp):
	logger.info("开始使用UVR5模型{0}：{1}",model,id)
	try:
		client1 = Client("http://localhost:7865/")
		# client1 = Client("http://localhost:9873/")
		result = client1.predict(
			model,  # str (Option from: []) in '模型' Dropdown component
			data,  # str in '输入待处理音频文件夹路径' Textbox component
			data_vocal,  # str in '指定输出主人声文件夹' Textbox component
			[],  # List[str] (List of filepath(s) or URL(s) to files) in '也可批量输入音频文件, 二选一, 优先读文件夹' File component
			data_temp,  # str in '指定输出非主人声文件夹' Textbox component
			10,  # float (numeric value between 0 and 20) in '人声提取激进程度' Slider component
			"wav",  # str in '导出文件格式' Radio component
			api_name="/uvr_convert"
		)


		print(result)
		# wirteMgs(id, "使用UVR5分离人声去除混响成功=="+result, 1)
	except Exception as e:
		traceback.print_exc()
		# wirteMgs(id,"使用UVR5分离人声去除混响失败=="+str(e),3)

#数据处理
def train_preprocess(id,project,data_row):
	logger.info("开始数据处理{0}：{1}",  id,data_row)
	client = Client("http://localhost:7865/")
	result2 = client.predict(
					data_row,	# str  in '输入训练文件夹路径' Textbox component
					project,	# str  in '输入实验名' Textbox component
					"40k",	# str  in '目标采样率' Radio component
					18,	# float (numeric value between 0 and 28) in '提取音高和处理数据使用的CPU进程数' Slider component
					api_name="/train_preprocess"
	)
	print(result2)

#特征提取
def train_extract_f0_feature(id,project):
	logger.info("开始特征提取{0}：{1}", id, project)
	client = Client("http://localhost:7865/")
	result3 = client.predict(
		"0",  # str  in '以-分隔输入使用的卡号, 例如   0-1-2   使用卡0和卡1和卡2' Textbox component
		19,  # float (numeric value between 0 and 28) in '提取音高和处理数据使用的CPU进程数' Slider component
		"rmvpe_gpu",
		# str  in '选择音高提取算法:输入歌声可用pm提速,高质量语音但CPU差可用dio提速,harvest质量更好但慢,rmvpe效果最好且微吃CPU/GPU' Radio component
		"true",  # str  in '模型是否带音高指导(唱歌一定要, 语音可以不要)' Radio component
		project,  # str  in '输入实验名' Textbox component
		"v2",  # str  in '版本' Radio component
		"0-0",  # str  in 'rmvpe卡号配置：以-分隔输入使用的不同进程卡号,例如0-0-1使用在卡0上跑2个进程并在卡1上跑1个进程' Textbox component
		api_name="/train_extract_f0_feature"
	)
	print(result3)

#训练模型
def trainModel(id,project,data_row):
	logger.info("开始训练模型{0}：{1}", id, project)
	client = Client("http://localhost:7865/")
	result = client.predict(
		project,  # str  in '输入实验名' Textbox component
		"40k",  # str  in '目标采样率' Radio component
		"true",  # str  in '模型是否带音高指导(唱歌一定要, 语音可以不要)' Radio component
		data_row,  # str  in '输入训练文件夹路径' Textbox component
		0,  # float (numeric value between 0 and 4) in '请指定说话人id' Slider component
		18,  # float (numeric value between 0 and 28) in '提取音高和处理数据使用的CPU进程数' Slider component
		"rmvpe",  # str  in '选择音高提取算法:输入歌声可用pm提速,高质量语音但CPU差可用dio提速,harvest质量更好但慢,rmvpe效果最好且微吃CPU/GPU' Radio component
		5,  # float (numeric value between 1 and 50) in '保存频率save_every_epoch' Slider component
		20,  # float (numeric value between 2 and 1000) in '总训练轮数total_epoch' Slider component
		12,  # float (numeric value between 1 and 40) in '每张显卡的batch_size' Slider component
		"是",  # str  in '是否仅保存最新的ckpt文件以节省硬盘空间' Radio component
		"assets/pretrained_v2/f0G40k.pth",  # str  in '加载预训练底模G路径' Textbox component
		"assets/pretrained_v2/f0D40k.pth",  # str  in '加载预训练底模D路径' Textbox component
		"0",  # str  in '以-分隔输入使用的卡号, 例如   0-1-2   使用卡0和卡1和卡2' Textbox component
		"否",  # str  in '是否缓存所有训练集至显存. 10min以下小数据可缓存以加速训练, 大数据缓存会炸显存也加不了多少速' Radio component
		"否",  # str  in '是否在每次保存时间点将最终小模型保存至weights文件夹' Radio component
		"v2",  # str  in '版本' Radio component
		"0-0",  # str  in 'rmvpe卡号配置：以-分隔输入使用的不同进程卡号,例如0-0-1使用在卡0上跑2个进程并在卡1上跑1个进程' Textbox component
		api_name="/train_start_all"
	)
	print(result)

#上传模型到指定服务器
def transfer(id,project_name):
	logger.info("上传模型到指定服务器："+project_name)
	try:
		list = []
		ssh = paramiko.SSHClient()
		ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
		ssh.connect(hostname='114.132.171.69', username='root', password='miaomiao2024!@#1')
		#传输模型
		model_source_file = model_path+project_name+'.pth'
		target_path = target_folder+project_name
		model_target_file = target_path+'/'+project_name+'.pth'
		list.append(target_real_folder+project_name+'/'+project_name+'.pth')
		#传输index文件
		index_source_file_name=""
		index_source_file_prefix =  'added_IVF'
		index_path_dir = index_path + project_name
		for filename in os.listdir(index_path_dir):
			if filename.startswith(index_source_file_prefix):
				index_source_file_name = filename
				print(filename)
		index_source_file = index_path_dir +'/'+index_source_file_name
		index_target_file = target_path +'/'+index_source_file_name
		list.append(target_real_folder+project_name+'/'+index_source_file_name)
		# 使用SFTP传输文件
		sftp = ssh.open_sftp()
		try:
			sftp.chdir(target_folder+project_name) # Test if remote_path exists
		except IOError:
			sftp.mkdir(target_folder+project_name)

		sftp.put(model_source_file, model_target_file)
		sftp.put(index_source_file, index_target_file)
		sftp.close()
		# 关闭SSH连接
		ssh.close()
		print("上传模型到指定服务器结束")
		# wirteMgs(id, "上传模型到指定服务器成功==", 1)
		return list
	except Exception as e:
		traceback.print_exc()
		# wirteMgs(id, "上传模型到指定服务器失败=="+str(e), 3)

# 写入数据库配置文件
def wirteDateBase(id,project_name, common_spk,common_name,list):
	print("写入数据库配置文件："+project_name)
	try:
	#查找参考文件对应音频
		data = ModelConfig(model_id=1000000+id,
						   speaker_zh=common_spk,
						   speaker=common_name,
						   model_path=list[0],
						   index_path=list[1],
						   language='zh',
						   device='cuda',
						   state=1,
						   length=1)
		db = SessionLocal()
		db.add(data)
		db.query(CustomizeTone).filter_by(id=id).update({"msg": "训练成功", "status": 2,"model_id":1000000+id})
		db.commit()
		db.close()
		print("写入数据库配置文件结束")
	except Exception as e:
		traceback.print_exc()


def wirteMgs(id,msg,status):
	db = SessionLocal()
	try:
		db.query(CustomizeTone).filter_by(id=id).update({"msg": msg, "status": status})
	except Exception as e:
		traceback.print_exc()
	finally:
		db.commit()
		db.close()

def download_file(url):
    response = requests.get(url, stream=True)
    response.raise_for_status()
    return response.content
def save_file(file_data, file_path):
    with open(file_path, 'wb') as file:
        file.write(file_data)


if __name__ == "__main__":
	# down_videos(5, "https://miaomiao-xuanyin.oss-cn-beijing.aliyuncs.com/test/sunce.mp3", "E:\\train_5_1")
	# uvr5_vocal(5, "HP3_all_vocals","E:\\train_5_1","E:\\train_5_1_temp","E:\\train_5_1_1")
	# uvr5_vocal(5, "HP5_only_main_vocal", "E:\\train_5_1_1","E:\\train_5_1_2", "E:\\train_5_1_temp")
	# uvr5_vocal(5, "onnx_dereverb_By_FoxJoy", "E:\\train_5_1_2", "E:\\train_5_1_temp","E:\\train_5_1_3")
	# uvr5_vocal(5, "VR-DeEchoDeReverb", "E:\\train_5_1_3", "E:\\train_5_1_final","E:\\train_5_1_temp",)
	# train_preprocess(5,"train_5","E:\\train_5_1_final")
	# train_extract_f0_feature(5,"train_5")
	# trainModel(5,"train_5","E:\\train_5_1_final")
	# list = transfer(5,"train_5")
	# wirteDateBase(5, "train_5",  "test", "test",list)  # 写入数据库
	logger.warning("本地服务，请勿将服务端口暴露于外网:8088")
	uvicorn.run(app, host='0.0.0.0', port=8088, workers=1)